# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Space updater"""

from typing import Union, List, Tuple, Optional
from numpy import ndarray

# convert Mindsore to Pytorch
import torch   
from torch import Tensor
from torch.nn import ModuleList, Parameter
from torch.optim import Optimizer


from ..system import Molecule
from ..control import Controller, Integrator, get_integrator,Thermostat, get_thermostat, Barostat, get_barostat
from ..function import get_ms_array, get_arguments, get_initializer, get_integer, GLOBAL_DEVICE_SPONGE

class Updater(Optimizer):
    
    def __init__(self,
                 system: Molecule,
                 controller: Union[Controller, List[Controller]] = None,
                 time_step: float = 1e-3,
                 velocity: Union[Tensor, ndarray, List[float]] = None,
                 weight_decay: float = 0.0,
                 loss_scale: float = 1.0,
                 foreach: Optional[bool] = None,
                 maximize: bool = False, capturable: bool = False,
                 differentiable: bool = False, fused: Optional[bool] = None,
                 **kwargs
                 ):

        self._kwargs = get_arguments(locals(), kwargs)
        self._kwargs.pop('velocity')

        self.device = GLOBAL_DEVICE_SPONGE()
        self.time_step = get_ms_array(time_step, dtype=torch.float32, device=self.device)

        self.system = system
        self.coordinate = self.system.coordinate.to(self.device)
        self.pbc_box = self.system.pbc_box.to(self.device) if self.system.pbc_box is not None else None

        # (B,A)
        self.atom_mass = self.system.atom_mass.to(self.device)
        self.inv_mass = self.system.inv_mass.to(self.device)
        # (B,A,1)
        self._atom_mass = self.atom_mass.unsqueeze(-1)
        self._inv_mass = self.inv_mass.unsqueeze(-1)

        self.num_walker = system.num_walker
        self.num_atoms = system.num_atoms
        self.dimension = system.dimension

        self.units = self.system.units
        self.kinetic_unit_scale = get_ms_array(self.units.kinetic_ref, dtype=torch.float32, device=self.device)

        if velocity is None:
            self.velocity = torch.zeros_like(self.coordinate, device=self.device)
        else:
            velocity = get_ms_array(velocity, torch.float32, device=self.device)
            if velocity.ndim == 2:
                velocity = velocity.unsqueeze(0)
            if velocity.shape != self.coordinate.shape:
                raise ValueError(f'The shape of velocity {velocity.shape} must be equal to '
                                 f'the shape of coordinate {self.coordinate.shape}!')
            self.velocity = get_ms_array(velocity, device=self.device)

        self.num_constraints = 0
        self.num_controller = 0
        self.controller: List[Controller] = None
        if controller is not None:
            if isinstance(controller, Controller):
                self.num_controller = 1
                controller = [controller]
            elif isinstance(controller, list):
                self.num_controller = len(controller)
            else:
                raise TypeError(f'The type of "controller" must be Controller or list but got: {type(controller)}')

            self.controller = ModuleList(controller)
            for i in range(self.num_controller):
                self.num_constraints += self.controller[i].num_constraints

        self.sys_dofs = system.degrees_of_freedom
        self.degrees_of_freedom = 0
        self.set_degrees_of_freedom(self.sys_dofs - self.num_constraints)

        self.kinetics = None
        self.temperature = None
        if self.velocity is not None:
            kinetics = self.get_kinetics(self.velocity)
            temperature = self.get_temperature(kinetics)
            # (B,D)
            self.kinetics = get_ms_array(kinetics, device=self.device)
            # (B)
            self.temperature = get_ms_array(temperature, device=self.device)

        self.virial = None
        self.pressure = None
        if self.pbc_box is not None:
            # (B,D)
            self.virial = torch.zeros(self.num_walker, self.dimension, device=self.device, requires_grad=True)
            self.pressure = torch.zeros(self.num_walker, self.dimension, device=self.device, requires_grad=True)

        self.step_num = get_ms_array(0, dtype=torch.int32, device=self.device)

        if controller is not None:
            for i in range(self.num_controller):
                self.controller[i].set_time_step(self.time_step)
        params = [system.coordinate]
        defaults = dict(time_step=time_step,maximize=maximize, foreach=foreach, capturable=capturable,
                        differentiable=differentiable, fused=fused)
        super().__init__(params=params, defaults=defaults)
        

    @property
    def boltzmann(self) -> float:

        return self.units.boltzmann

    @property
    def press_unit_scale(self) -> float:

        return self.units.pressure_ref

    def set_step(self, step: int = 0):

        step = get_ms_array(step, dtype=torch.int32, device=self.device)
        return self

    def set_degrees_of_freedom(self, dofs: int):
        self.degrees_of_freedom = get_integer(dofs)
        self.num_constraints = self.sys_dofs - self.degrees_of_freedom
        for i in range(self.num_controller):
            self.controller[i].set_degrees_of_freedom(self.degrees_of_freedom)
        return self

    def update_coordinate(self, coordinate: Tensor, success: bool = True) -> bool:

        self.coordinate.data.copy_(coordinate)
        return success

    def update_pbc_box(self, pbc_box: Tensor, success: bool = True) -> bool:

        if self.pbc_box is None:
            return success
        self.pbc_box.data.copy_(pbc_box)
        return success

    def update_velocity(self, velocity: Tensor, success: bool = True) -> bool:

        self.velocity.data.copy_(velocity)
        return success

    def update_kinetics(self, kinetics: Tensor, success: bool = True) -> bool:

        if self.kinetics is None:
            return success
        self.kinetics.data.copy_(kinetics)
        return success

    def update_temperature(self, temperature: Tensor, success: bool = True) -> bool:

        if self.temperature is None:
            return success
        self.temperature.data.copy_(temperature)
        return success

    def update_virial(self, virial: Tensor, success: bool = True) -> bool:

        if self.pbc_box is None:
            return success
        self.virial.data.copy_(virial)
        return success

    def update_pressure(self, pressure: Tensor, success: bool = True) -> bool:

        if self.pbc_box is None:
            return success
        self.pressure.data.copy_(pressure)
        return success

    def get_velocity(self) -> Tensor:

        if self.velocity is None:
            return None
        return self.velocity.data

    def get_kinetics(self, velocity: Tensor) -> Tensor:

        # (B,A,D)
        kinetics = 0.5 * self._atom_mass * velocity**2
        # (B,D) <- (B,A,D)
        kinetics = torch.sum(kinetics, -2)
        return kinetics * self.kinetic_unit_scale

    def get_temperature(self, kinetics: Tensor = None) -> Tensor:

        # (B) <- (B,D)
        kinetics = torch.sum(kinetics, -1)
        return 2 * kinetics / self.degrees_of_freedom / self.boltzmann

    def get_pressure(self, kinetics: Tensor, virial: Tensor, pbc_box: Tensor) -> Tensor:

        if self.pbc_box is None:
            return None
        # (B,D) = ((B,D) - (B, D)) / (B,1)
        volume = torch.prod(pbc_box, -1, keepdim=True)
        pressure = 2 * (kinetics - virial) / volume
        return pressure * self.press_unit_scale

    def next_step(self, success: bool = True) -> bool:

        self.step_num.data.copy_(self.step_num + 1)
        return success

    def step(self, energy: Tensor, force: Tensor, virial: Tensor = None):
        """update the parameters of system"""

        coordinate = self.coordinate
        velocity = self.velocity
        kinetics = self.kinetics
        pbc_box = self.pbc_box

        step = self.step_num.data
        if self.controller is not None:
            for i in range(self.num_controller):
                coordinate, velocity, force, energy, kinetics, virial, pbc_box = \
                    self.controller[i](coordinate, velocity, force, energy, kinetics, virial, pbc_box, step)

        temperature = self.get_temperature(kinetics)
        pressure = self.get_pressure(kinetics, virial, pbc_box)

        success = True
        success = self.update_coordinate(coordinate, success)
        success = self.update_velocity(velocity, success)
        success = self.update_pbc_box(pbc_box, success)
        success = self.update_kinetics(kinetics, success)
        success = self.update_temperature(temperature, success)
        success = self.update_virial(virial, success)
        success = self.update_pressure(pressure, success)

        return self.next_step(success)
